import time
import re
import json
import threading
import functools
from openai import OpenAI
import tiktoken

enc = tiktoken.encoding_for_model("gpt-4")

zhipu_key = '8c916e6b87d0498d9b65bf8134353904.1FdeqviAGsIwCSES'

client = OpenAI(
    api_key=zhipu_key,
    base_url="https://open.bigmodel.cn/api/paas/v4/"
)


class TimeoutException(Exception):
    pass


def timeout(seconds):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            result = [TimeoutException(f"Function '{func.__name__}' timed out after {seconds} seconds")]

            def target():
                try:
                    result[0] = func(*args, **kwargs)
                except Exception as e:
                    result[0] = e

            thread = threading.Thread(target=target)
            thread.start()
            thread.join(seconds)
            if thread.is_alive():
                raise result[0]
            return result[0]

        return wrapper

    return decorator


import functools
import traceback


def try_n_times(n):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for i in range(n):
                try:
                    result = func(*args, **kwargs)
                    return result
                except Exception as e:
                    traceback.print_exc()
                    if i == n - 1:  # 当达到最大尝试次数时，抛出异常
                        raise e
                    print(f"Attempt {i + 1} failed. Retrying...")
                    time.sleep(20)  # 可选，等待1秒后重试

        return wrapper

    return decorator


def super_eval(json_str, try_num=0):
    if try_num > 3:
        return 'json格式错误'
    json_str = json_str.replace('：', ':')
    try:
        all_json = re.findall('```json(.*?)```', json_str, re.DOTALL)
        if all_json:
            try:
                return json.loads(all_json[-1])
            except:
                import traceback
                traceback.print_exc()
                return eval(all_json[-1])

        if '```json' in json_str:
            json_str = json_str.replace('```json', '')
        json_str = json_str.replace('```', '')
        try:
            return eval(json_str)
        except:
            return json.loads(json_str)
    except:
        import traceback
        traceback.print_exc()
        text = llm(f"输出以下内容的json部分并修复成正确格式备注仅仅输出最后的json:```{json_str}```")
        try_num += 1
        return super_eval(text, try_num)


@try_n_times(3)
def llmglm(content, print_str=True, max_input=60000, max_rounds=30):
    if isinstance(content, str):
        messages = [{'role': 'user', 'content': content}]
    elif isinstance(content, list):
        if len(content) > max_rounds:
            raise ValueError('对话轮数过长')
        # print(content[-1])
        messages = content
    else:
        raise ValueError('输入格式错误')
    if len(enc.encode(str(content))) > max_input:
        raise ValueError('输入Token过长')
    completion = client.chat.completions.create(
        model="glm-4.5",
        messages=messages,
        temperature=0.2,
        # max_tokens=5000,
        timeout=1000,
        tools=[{'type': 'web_search', 'web_search': {"enable": False, 'type': 'web_search'}}],
    )
    answer = '<think>' + completion.choices[0].message.reasoning_content + '</think>\n' + completion.choices[
        0].message.content

    if isinstance(content, list):
        content.append({"role": "assistant", "content": answer})
    return answer


deepseek_key = 'sk-327348b052384532824142a19d5174f3'
deepseek_url = 'https://api.deepseek.com'

deepseek_client = OpenAI(
    api_key=deepseek_key,
    base_url=deepseek_url
)


# @try_n_times(3)
def llmds(content, print_str=True, max_input=50000, max_rounds=300):
    print('deep seek')
    if isinstance(content, str):
        messages = [{'role': 'user', 'content': content}]
    elif isinstance(content, list):
        if len(content) > max_rounds:
            raise ValueError('对话轮数过长')
        print(content[-1])
        messages = content
    else:
        raise ValueError('输入格式错误')

    completion = deepseek_client.chat.completions.create(
        model="deepseek-reasoner",
        # model="deepseek-chat",
        messages=messages,
        temperature=0.2)

    # answer = "<think>"+completion.choices[0].message.reasoning_content+"</think>" + completion.choices[0].message.content
    answer = completion.choices[0].message.content

    if isinstance(content, list):
        content.append({"role": "assistant", "content": answer})

    print(answer)
    print(completion.usage)
    return answer


llm = llmds

def llm_eval(content):
    return super_eval(llm(content))

if __name__ == '__main__':
    print(llm('你是谁'))
